#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
检查GPU和CUDA可用性
验证系统是否可以使用GPU进行模型训练
"""

import sys
import os

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, project_root)

def check_gpu_availability():
    """检查GPU和CUDA可用性"""
    print("[INFO] GPU和CUDA可用性检查")
    print("=" * 60)
    
    try:
        import torch
        print("[INFO] PyTorch版本:", torch.__version__)
        
        # 检查CUDA可用性
        cuda_available = torch.cuda.is_available()
        print(f"[INFO] CUDA可用性: {cuda_available}")
        
        if cuda_available:
            # 获取CUDA详细信息
            cuda_version = torch.version.cuda
            gpu_count = torch.cuda.device_count()
            current_device = torch.cuda.current_device()
            gpu_name = torch.cuda.get_device_name(current_device)
            
            print(f"[INFO] CUDA版本: {cuda_version}")
            print(f"[INFO] GPU数量: {gpu_count}")
            print(f"[INFO] 当前GPU: {current_device}")
            print(f"[INFO] GPU名称: {gpu_name}")
            
            # 显存信息
            gpu_memory = torch.cuda.get_device_properties(current_device).total_memory
            gpu_memory_gb = gpu_memory / (1024**3)
            print(f"[INFO] GPU显存: {gpu_memory_gb:.2f} GB")
            
            # 当前显存使用情况
            allocated_memory = torch.cuda.memory_allocated() / (1024**3)
            cached_memory = torch.cuda.memory_reserved() / (1024**3)
            print(f"[INFO] 已分配显存: {allocated_memory:.2f} GB")
            print(f"[INFO] 缓存显存: {cached_memory:.2f} GB")
            
            return True
        else:
            print("[WARNING] CUDA不可用")
            print("💡 可能的原因:")
            print("  1. 未安装支持CUDA的PyTorch版本")
            print("  2. 显卡驱动问题")
            print("  3. CUDA版本不兼容")
            print("  4. 环境配置问题")
            return False
            
    except ImportError as e:
        print(f"[ERROR] 无法导入PyTorch: {e}")
        print("💡 请确保已正确安装PyTorch")
        return False
    except Exception as e:
        print(f"[ERROR] 检查过程中出错: {e}")
        return False

def check_pytorch_installation():
    """检查PyTorch安装详情"""
    print("\n🔍 PyTorch安装详情")
    print("=" * 60)
    
    try:
        import torch
        print("[INFO] PyTorch已安装")
        print(f"  版本: {torch.__version__}")
        print(f"  安装路径: {torch.__file__}")
        
        # 检查构建配置
        print(f"  构建信息: {torch.__config__.show()}")
        
        # 检查是否支持CUDA
        if torch.cuda.is_available():
            print("  CUDA支持: [INFO] 可用")
        else:
            print("  CUDA支持: [ERROR] 不可用")
            
    except Exception as e:
        print(f"[ERROR] 检查PyTorch安装详情失败: {e}")

def check_conda_environment():
    """检查conda环境"""
    print("\n🐍 Conda环境检查")
    print("=" * 60)
    
    try:
        # 检查conda环境变量
        conda_prefix = os.environ.get('CONDA_PREFIX')
        if conda_prefix:
            print(f"[INFO] Conda环境: {conda_prefix}")
        else:
            print("[WARNING] 未检测到Conda环境")
            
        # 检查Python路径
        print(f"🐍 Python路径: {sys.executable}")
        print(f"📦 Python版本: {sys.version}")
        
    except Exception as e:
        print(f"[ERROR] 检查conda环境失败: {e}")

def test_gpu_tensor_operations():
    """测试GPU张量操作"""
    print("\n🧪 GPU张量操作测试")
    print("=" * 60)
    
    try:
        import torch
        
        if torch.cuda.is_available():
            # 创建测试张量
            print("创建测试张量...")
            cpu_tensor = torch.randn(2, 3, 4)
            print(f"[INFO] CPU张量创建成功，形状: {cpu_tensor.shape}")
            
            # 移动到GPU
            print(" 将张量移动到GPU...")
            gpu_tensor = cpu_tensor.cuda()
            print(f"[INFO] GPU张量创建成功，形状: {gpu_tensor.shape}")
            print(f"📍 张量设备: {gpu_tensor.device}")
            
            # 执行简单计算
            print("🔢 执行GPU计算...")
            result = torch.matmul(gpu_tensor, gpu_tensor.transpose(-1, -2))
            print(f"[INFO] 矩阵乘法成功，结果形状: {result.shape}")
            print(f"📍 结果设备: {result.device}")
            
            print("[INFO] GPU测试完全成功！")
            return True
        else:
            print("[WARNING] CUDA不可用，跳过GPU测试")
            return False
            
    except Exception as e:
        print(f"[ERROR] GPU张量操作测试失败: {e}")
        return False

def check_plw_model_gpu_support():
    """检查排列5模型GPU支持"""
    print("\n🎯 排列5模型GPU支持检查")
    print("=" * 60)
    
    try:
        import torch
        from algorithms.plw_sequence_lstm import PLWSequenceLSTM
        
        # 创建小型模型用于测试
        print("[INFO] 创建PLW序列LSTM模型...")
        model = PLWSequenceLSTM(input_dim=5, hidden_dim=32, num_layers=1, dropout=0.1)
        print(f"[INFO] 模型创建成功")
        
        if torch.cuda.is_available():
            # 将模型移动到GPU
            print(" 将模型移动到GPU...")
            model = model.cuda()
            print(f"[INFO] 模型已移动到设备: {next(model.parameters()).device}")
            
            # 创建测试输入
            print("📥 创建测试输入...")
            batch_size, seq_len, input_dim = 2, 5, 5
            test_input = torch.randint(0, 10, (batch_size, seq_len, input_dim)).float().cuda()
            print(f"[INFO] 输入张量创建成功，设备: {test_input.device}")
            
            # 前向传播测试
            print("[INFO] 执行前向传播...")
            model.eval()
            with torch.no_grad():
                output = model(test_input)
                print(f"[INFO] 前向传播成功，输出形状: {output.shape}")
                print(f"📍 输出设备: {output.device}")
            
            print("[INFO] PLW模型GPU支持测试成功！")
            return True
        else:
            # CPU测试
            print("📥 创建测试输入(CPU)...")
            batch_size, seq_len, input_dim = 2, 5, 5
            test_input = torch.randint(0, 10, (batch_size, seq_len, input_dim)).float()
            print(f"[INFO] 输入张量创建成功，设备: {test_input.device}")
            
            # 前向传播测试
            print("[INFO] 执行前向传播(CPU)...")
            model.eval()
            with torch.no_grad():
                output = model(test_input)
                print(f"[INFO] 前向传播成功，输出形状: {output.shape}")
                print(f"📍 输出设备: {output.device}")
            
            print("[WARNING] CUDA不可用，使用CPU测试")
            return True
            
    except Exception as e:
        print(f"[ERROR] PLW模型GPU支持检查失败: {e}")
        import traceback
        print(traceback.format_exc())
        return False

if __name__ == "__main__":
    print("[INFO] 开始GPU和CUDA可用性检查")
    print("[INFO] 这将帮助诊断排列5模型训练使用CPU而非GPU的问题")
    print()
    
    # 检查GPU可用性
    gpu_available = check_gpu_availability()
    
    # 检查PyTorch安装
    check_pytorch_installation()
    
    # 检查conda环境
    check_conda_environment()
    
    # 测试GPU张量操作
    gpu_test_success = test_gpu_tensor_operations()
    
    # 检查PLW模型GPU支持
    plw_gpu_success = check_plw_model_gpu_support()
    
    print("\n" + "=" * 80)
    print("检查结果总结:")
    print(f"  GPU可用性: {'[INFO] 可用' if gpu_available else '[ERROR] 不可用'}")
    print(f"  GPU张量测试: {'[INFO] 通过' if gpu_test_success else '[ERROR] 失败'}")
    print(f"  PLW模型GPU支持: {'[INFO] 通过' if plw_gpu_success else '[ERROR] 失败'}")
    
    if gpu_available and gpu_test_success and plw_gpu_success:
        print("\n[INFO] GPU配置正常！")
        print("💡 排列5模型训练应该可以使用GPU")
        print("\n[INFO] 如果仍在使用CPU，请检查:")
        print("  1. 训练脚本中的设备配置")
        print("  2. 模型加载时的设备设置")
        print("  3. 环境变量配置")
    else:
        print("\n[WARNING] GPU配置存在问题")
        print("💡 建议:")
        print("  1. 重新安装支持CUDA的PyTorch")
        print("  2. 检查显卡驱动")
        print("  3. 验证CUDA版本兼容性")